import glob
import numpy as np
import torchio as tio
import logging
import os
import skimage.exposure as exposure
import nibabel as nib
import matplotlib.pyplot as plt


def resample_by_torchio(input_path, output_path):
    # 加载图像
    image = tio.ScalarImage(input_path)

    # 定义变换
    transform = tio.Compose(
        [
            tio.ToCanonical(),  # 将图像转换为标准方向
            tio.RescaleIntensity((0, 255)),  # 重新缩放强度到 [0, 255]
        ]
    )

    # 应用变换
    transformed_image = transform(image)

    # 获取图像数据并转换为 uint8
    img_data = transformed_image.data.numpy()[0]  # 形状 (x, y, z)
    img_data_uint8 = img_data.astype(np.uint8)  # 强制转换为 uint8

    # 使用 nibabel 保存为 NIfTI 格式
    nifti_img = nib.Nifti1Image(img_data_uint8, affine=transformed_image.affine)
    nifti_img.header.set_data_dtype(np.uint8)  # 显式指定数据类型为 uint8
    nib.save(nifti_img, output_path)

    # 调试：检查保存后的范围
    saved_image = tio.ScalarImage(output_path)
    print(f"Saved range: min={saved_image.data.min()}, max={saved_image.data.max()}")


def resample_unite2(dir_list=["./proposed"], target_dir="./processed"):
    """
    将多个目录中的图像统一重采样到按照规则自动创建的目标目录。
    Params:
        dir_list: 需要处理的目录列表
        target_dir: 统一的目标目录
    """
    for dir in dir_list:
        if os.path.exists(target_dir):
            pass
        else:
            os.makedirs(target_dir, exist_ok=True)
            logging.info(f"Created target directory: {target_dir}")
        input_files = glob.glob(os.path.join(dir, "**", "*.nii.gz"), recursive=True)
        # input_files = [
        #     file
        #     for file in input_files
        #     if os.path.basename(file).startswith(f"{subject_id}")
        # ]
        logging.info(f"Processing files{input_files}")
        for img_path in input_files:
            output_path = os.path.basename(img_path).replace(
                ".nii.gz", "_registered.nii.gz"
            )
            output_path = os.path.join(target_dir, output_path)
            resample_by_torchio(img_path, output_path)


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
    )
    logger = logging.getLogger(__name__)
    dir_list = [
        "../prediction/CHEN-REN-GENG_tra/proposed/cubic__fix_ds_6x4",
    ]
    target_dir = "../prediction/CHEN-REN-GENG_tra/processed"
    resample_unite2(
        dir_list,
        target_dir=target_dir,
    )
